from __future__ import division

import logging, unittest, time

import numpy.testing, scipy.integrate, scipy.constants

import common, wlc, kinetics_opt

logger = logging.getLogger("kinetics_opt_test")
logger.addHandler(logging.StreamHandler())
logger.setLevel(logging.INFO)

def slow_integrals(
    t, p, dx, upper_bound_units, newtons_per_unit, peak_integral_params_list):
    kbt_over_p = scipy.constants.k * t / p
    kbt_over_dx = scipy.constants.k * t / dx
    n = kinetics_opt.PeakIntegralParams.sum_counts(peak_integral_params_list)
    a = numpy.empty(n, dtype=kinetics_opt.bestfloat)
    i = 0
    for pip in peak_integral_params_list:
        j = i + pip.count()
        def func(z):
            return numpy.exp(z / kbt_over_dx) * wlc.f_inverse_derived_approximation(
                kbt_over_p, pip.kl, pip.v_over_l, z)
        for k in range(i, j):
            a[k] = scipy.integrate.quad(
                func,
                pip.y_units[k - i] * newtons_per_unit,
                pip.y_units[k - i + 1] * newtons_per_unit)[0]
        a[i:j].cumsum(out=a[i:j])
        i = j
    return a

def check_sum_y_of_w_is_one(s, f):
    for upper_bound_units in range(s + 1, s + 101, 7):
        logger.debug("s=%d upper_bound_units=%d", s, upper_bound_units)
        y = s
        p = 0
        while True and y >= 3:
            assert y != upper_bound_units
            basic_peak = kinetics_opt.BasicPeak(kinetics_opt.Protocol(1, 2, 3, 5), y)
            logger.debug(
                "y=%d upper_bound_units=%d s=%d basic_peak=%s f=%s",
                y, upper_bound_units, s, basic_peak, f)
            actual = f(basic_peak, 2, 3, upper_bound_units)
            filtered_weights = actual.weights[actual.y_end_units == s]
            common.check_le(len(filtered_weights), 1)
            if len(filtered_weights):
                p += filtered_weights.sum()
                y += ((y >= s) << 1) - 1
                if y == upper_bound_units:
                    y = s - 1
            else:
                common.check_not_equal(y, s, (f, s, y, actual))
                if y > s:
                    y = s - 1
                else:
                    break
        numpy.testing.assert_allclose(p, 1)


class KineticsOptTest(unittest.TestCase):

    def test_interval_covering(self):
        actual = kinetics_opt.interval_covering([(10, 20)])
        expected = ((10, 20), (set([0]), ))
        self.assertEqual(expected, actual)
        actual = kinetics_opt.interval_covering([(10, 20), (10, 20)])
        expected = ((10, 20), (set([0, 1]), ))
        self.assertEqual(expected, actual)
        actual = kinetics_opt.interval_covering([(30, 40), (10, 20)])
        expected = ((10, 20, 30, 40), (set([1]), set(), set([0])))
        self.assertEqual(expected, actual)
        actual = kinetics_opt.interval_covering([(20, 30), (10, 20)])
        expected = ((10, 20, 30), (set([1]), set([0])))
        self.assertEqual(expected, actual)
        actual = kinetics_opt.interval_covering([(10, 30), (10, 20)])
        expected = ((10, 20, 30), (set([0, 1]), set([0])))
        self.assertEqual(expected, actual)
        actual = kinetics_opt.interval_covering([(10, 30), (20, 40)])
        expected = ((10, 20, 30, 40), (set([0]), set([0, 1]), set([1])))
        self.assertEqual(expected, actual)

    def test_simple_integrals(self):
        actual = kinetics_opt.integrals(
            1, 1, 1, 1000, 1, [kinetics_opt.PeakIntegralParams(1, 1, [1])])
        numpy.testing.assert_equal(actual, [])
        actual = kinetics_opt.integrals(1, 1, 1, 1000, 1, [])
        numpy.testing.assert_equal(actual, [])

    def test_one_integral(self):
        t = 301.15
        p = wlc.default_p
        dx = 0.165e-9
        upper_bound_units = 1000
        newtons_per_unit = 1e-12
        v = 2180e-9
        l = 80e-9
        k = 0.06
        peak_integral_params_list = [
            kinetics_opt.PeakIntegralParams(v / l, k * l, [0, 200])]
        actual = kinetics_opt.integrals(
            t, p, dx, upper_bound_units, newtons_per_unit, peak_integral_params_list)
        expected = slow_integrals(
            t, p, dx, upper_bound_units, newtons_per_unit, peak_integral_params_list)
        numpy.testing.assert_allclose(actual, expected, rtol=3e-5)

    def test_integral_for_one_peak_with_two_forces(self):
        t = 301.15
        p = wlc.default_p
        dx = 0.165e-9
        upper_bound_units = 1000
        newtons_per_unit = 1e-12
        v = 2180e-9
        l = 80e-9
        k = 0.06
        peak_integral_params_list = [
            kinetics_opt.PeakIntegralParams(v / l, k * l, [0, 200, 300])]
        actual = kinetics_opt.integrals(
            t, p, dx, upper_bound_units, newtons_per_unit, peak_integral_params_list)
        expected = slow_integrals(
            t, p, dx, upper_bound_units, newtons_per_unit, peak_integral_params_list)
        numpy.testing.assert_allclose(actual, expected, rtol=3e-5)

    def test_two_integrals(self):
        t = 301.15
        p = wlc.default_p
        dx = 0.165e-9
        upper_bound_units = 1000
        newtons_per_unit = 1e-12
        v0, v1 = 50.1e-9, 2180e-9
        l0, l1 = 30e-9, 80e-9
        k0, k1 = 0.06, 0.12
        peak_integral_params_list = [
            kinetics_opt.PeakIntegralParams(v0 / l0, k0 * l0, [0, 200]),
            kinetics_opt.PeakIntegralParams(v1 / l1, k1 * l1, [0, 100]),
            ]
        actual = kinetics_opt.integrals(
            t, p, dx, upper_bound_units, newtons_per_unit, peak_integral_params_list)
        expected = slow_integrals(
            t, p, dx, upper_bound_units, newtons_per_unit, peak_integral_params_list)
        numpy.testing.assert_allclose(actual, expected, rtol=9e-4)

    def test_normal_from_basic_peak(self):
        upper_bound_units = 100
        protocol = kinetics_opt.Protocol(1, 2, 3, 5)
        basic_peak = kinetics_opt.BasicPeak(protocol, 44)
        actual = kinetics_opt.DiffusePeak.normal_from_basic_peak(
            basic_peak, 0, 0, upper_bound_units)
        expected = kinetics_opt.DiffusePeak(protocol, [1], [44])
        self.assertEqual(actual, expected)
        basic_peak = kinetics_opt.BasicPeak(protocol, 3)
        actual = kinetics_opt.DiffusePeak.normal_from_basic_peak(
            basic_peak, 0, 0, upper_bound_units)
        expected = kinetics_opt.DiffusePeak(protocol, [1], [3])
        self.assertEqual(actual, expected)
        basic_peak = kinetics_opt.BasicPeak(protocol, 3)
        actual = kinetics_opt.DiffusePeak.normal_from_basic_peak(
            basic_peak, 1e10, 1, upper_bound_units)
        expected = kinetics_opt.DiffusePeak(protocol, [1 / 2, 1 / 3], [3, 4])
        self.assertTrue(actual.allclose(expected), (actual, expected))
        basic_peak = kinetics_opt.BasicPeak(protocol, 4)
        actual = kinetics_opt.DiffusePeak.normal_from_basic_peak(
            basic_peak, 1, 1, upper_bound_units)
        h, l = scipy.stats.norm.pdf(0), scipy.stats.norm.pdf(1)
        h, l = h / (h + 2 * l), l / (h + 2 * l)
        expected = kinetics_opt.DiffusePeak(protocol, [l / (h + l), h, l], [3, 4, 5])
        self.assertTrue(actual.allclose(expected))

    def test_truncated_normal_from_basic_peak(self):
        upper_bound_units = 100
        protocol = kinetics_opt.Protocol(1, 2, 3, 5)
        basic_peak = kinetics_opt.BasicPeak(protocol, 44)
        actual = kinetics_opt.DiffusePeak.truncated_normal_from_basic_peak(
            basic_peak, 0, 0, upper_bound_units)
        expected = kinetics_opt.DiffusePeak(protocol, [1], [44])
        self.assertEqual(actual, expected)
        basic_peak = kinetics_opt.BasicPeak(protocol, 3)
        actual = kinetics_opt.DiffusePeak.truncated_normal_from_basic_peak(
            basic_peak, 0, 0, upper_bound_units)
        expected = kinetics_opt.DiffusePeak(protocol, [1], [3])
        self.assertEqual(actual, expected)
        basic_peak = kinetics_opt.BasicPeak(protocol, 4)
        actual = kinetics_opt.DiffusePeak.truncated_normal_from_basic_peak(
            basic_peak, 1e10, 1, upper_bound_units)
        expected = kinetics_opt.DiffusePeak(protocol, [1 / 2, 1 / 2], [3, 4])
        actual = kinetics_opt.DiffusePeak.truncated_normal_from_basic_peak(
            basic_peak, 1, 1, upper_bound_units)
        h, l = scipy.stats.norm.pdf(0), scipy.stats.norm.pdf(1)
        expected = kinetics_opt.DiffusePeak(
            protocol, [l / (h + l), h / (h + l)], [3, 4])
        self.assertTrue(actual.allclose(expected))

    def test_sum_y_of_w_is_one(self):
        for f in (
            kinetics_opt.DiffusePeak.normal_from_basic_peak,
            kinetics_opt.DiffusePeak.truncated_normal_from_basic_peak,
            ):
            for s in (3, 4, 5, 1000):
                check_sum_y_of_w_is_one(s, f)

    def test_full_partial_sums(self):
        kbt_over_dx = 4
        newtons_per_unit = 2
        h = numpy.array([999])
        y_start_list = [0, 0, 0]
        actual = kinetics_opt.full_partial_sums(
            kbt_over_dx, newtons_per_unit, h, y_start_list)
        numpy.testing.assert_array_equal(
            actual, numpy.zeros((3, 1, 2)))
        h = numpy.zeros(2)  # TODO: make it non-zero to check also diff1
        y_start_list = [1, 1, 0]
        actual = kinetics_opt.full_partial_sums(
            kbt_over_dx, newtons_per_unit, h, y_start_list)
        diff0 = numpy.exp(newtons_per_unit * 1 / kbt_over_dx) - numpy.exp(
            newtons_per_unit * 0 / kbt_over_dx)
        expected = numpy.array(
                [[[numpy.inf, numpy.inf], [0, 0]],
                 [[numpy.inf, numpy.inf], [0, 0]],
                 [[0, 0], [diff0, 0]]])
        numpy.testing.assert_allclose(actual, expected)
        h = numpy.zeros(3)
        y_start_list = [1, 1, 0]
        actual = kinetics_opt.full_partial_sums(
            kbt_over_dx, newtons_per_unit, h, y_start_list)
        diff0_a = numpy.exp(newtons_per_unit * 1 / kbt_over_dx) - numpy.exp(
            newtons_per_unit * 0 / kbt_over_dx)
        diff0_b = numpy.exp(newtons_per_unit * 2 / kbt_over_dx) - numpy.exp(
            newtons_per_unit * 1 / kbt_over_dx)
        expected = numpy.array(
                [[[numpy.inf, numpy.inf], [0, 0], [diff0_b, 0]],
                 [[numpy.inf, numpy.inf], [0, 0], [diff0_b, 0]],
                 [[0, 0], [diff0_a, 0], [diff0_a + diff0_b, 0]]])
        numpy.testing.assert_allclose(actual, expected)

    def test_full_log_pdfs(self):
        t = 301.15
        p = wlc.default_p
        dx = 0.165e-9
        k0 = 0.039
        upper_bound_units = 40
        newtons_per_unit = 10e-12
        v = 400e-9
        l = 80e-9
        k = 0.06
        m = 3
        for y_start_units in range(upper_bound_units):
            protocol = kinetics_opt.Protocol(v / l, k * l, y_start_units, m)
            actual = kinetics_opt.full_log_pdfs(
                t, p, upper_bound_units, newtons_per_unit, [protocol], dx, k0)
            basic_peaks = [
                kinetics_opt.BasicPeak(protocol, j)
                for j in range(y_start_units, upper_bound_units)]
            lls = kinetics_opt.log_likelihoods(
                t, p, dx, k0, upper_bound_units, newtons_per_unit, basic_peaks)
            expected = numpy.concatenate(
                [[-numpy.inf] * y_start_units, lls])[numpy.newaxis, :]
            numpy.testing.assert_allclose(actual, expected)

    def test_full_log_pdfs_diffuse(self):
        t = 301.15
        p = wlc.default_p
        dx = 0.165e-9
        k0 = 0.039
        upper_bound_units = 10
        newtons_per_unit = 10e-12
        v = 400e-9
        l = 80e-9
        k = 0.06
        m = 3
        sigma_units = 2
        hsu = 6
        for y_start_units in range(upper_bound_units):
            diffuse_protocol = kinetics_opt.DiffuseProtocol(
                kinetics_opt.Protocol(v / l, k * l, y_start_units, m),
                sigma_units, hsu)
            for f in (
                kinetics_opt.DiffusePeak.normal_from_basic_peak,
                kinetics_opt.DiffusePeak.truncated_normal_from_basic_peak,
                ):
                actual = kinetics_opt.full_log_pdfs_diffuse(
                    t, p, upper_bound_units, newtons_per_unit, [diffuse_protocol],
                    dx, k0, f)
                basic_peaks = [
                    kinetics_opt.BasicPeak(diffuse_protocol.protocol, j)
                    for j in range(y_start_units, upper_bound_units)]
                diffuse_peaks = [
                    f(basic_peak, sigma_units, hsu, upper_bound_units)
                    for basic_peak in basic_peaks]
                lls = kinetics_opt.log_likelihoods_diffuse(
                    t, p, dx, k0, upper_bound_units, newtons_per_unit, diffuse_peaks)
                expected = numpy.concatenate(
                    [[-numpy.inf] * y_start_units, lls])[numpy.newaxis, :]
                numpy.testing.assert_allclose(actual, expected)

    def test_compute_probabilities(self):
        log_pdfs = numpy.log(numpy.array([[1, 3, 2], [2, 1, 1]]))
        actual = kinetics_opt.compute_probabilities(log_pdfs)
        expected = numpy.array([[1/6, 3/6, 2/6], [2/4, 1/4, 1/4]])
        numpy.testing.assert_allclose(actual, expected)
        
    def test_estimate_shared_dx_k0_for_one_dataset_is_same_as_isolated(self):
        t = 301.15
        p = wlc.default_p
        dx = 0.165e-9
        k0 = 0.039
        upper_bound_units = 400
        newtons_per_unit = 1e-12
        v = 400e-9
        l = 80e-9
        k = 0.06
        m = 3
        y_start_units = int(numpy.around(10e-12 / newtons_per_unit))
        protocol = kinetics_opt.Protocol(v / l, k * l, y_start_units, m)
        n = 10
        basic_peaks = [
            kinetics_opt.BasicPeak(protocol, protocol.y_start_units * (i + 1))
            for i in range(n)]
        mle_dx, mle_k0, mle_lls = kinetics_opt.estimate_dx_k0(
            t, p, upper_bound_units, newtons_per_unit, basic_peaks,
            0.01e-9, 1e-9, False)
        actual_dx, actual_k0_lls_list = kinetics_opt.estimate_shared_dx_k0(
            t, p, upper_bound_units, newtons_per_unit, [basic_peaks],
            0.01e-9, 1e-9, False)
        self.assertEqual(actual_dx, mle_dx)
        self.assertEqual(len(actual_k0_lls_list), 1)
        self.assertEqual(actual_k0_lls_list[0][0], mle_k0)
        numpy.testing.assert_array_equal(actual_k0_lls_list[0][1], mle_lls)
        
    def test_estimate_shared_dx_k0_diffuse_for_one_dataset_is_same_as_isolated(self):
        t = 301.15
        p = wlc.default_p
        dx = 0.165e-9
        k0 = 0.039
        upper_bound_units = 400
        newtons_per_unit = 1e-12
        v = 400e-9
        l = 80e-9
        k = 0.06
        m = 3
        sigma_units = 20
        hsu = 60
        y_start_units = int(numpy.around(10e-12 / newtons_per_unit))
        diffuse_protocol = kinetics_opt.DiffuseProtocol(
            kinetics_opt.Protocol(v / l, k * l, y_start_units, m),
            sigma_units, hsu)
        for factory in (
            kinetics_opt.DiffusePeak.normal_from_basic_peak,
            kinetics_opt.DiffusePeak.truncated_normal_from_basic_peak,
            ):
            n = 10
            diffuse_peaks = [
                factory(
                    kinetics_opt.BasicPeak(
                        diffuse_protocol.protocol,
                        diffuse_protocol.y_start_units() * (i + 1)),
                    diffuse_protocol.sigma_units,
                    diffuse_protocol.hsu, upper_bound_units)
                for i in range(n)]
            mle_dx, mle_k0, mle_lls = kinetics_opt.estimate_dx_k0_diffuse(
                t, p, upper_bound_units, newtons_per_unit, diffuse_peaks,
                0.01e-9, 1e-9,
                kinetics_opt.default_diffuse_k0_lower_bound,
                kinetics_opt.default_diffuse_k0_upper_bound)
            actual_dx, actual_k0_lls_list = kinetics_opt.estimate_shared_dx_k0_diffuse(
                t, p, upper_bound_units, newtons_per_unit, [diffuse_peaks],
                0.01e-9, 1e-9,
                kinetics_opt.default_diffuse_k0_lower_bound,
                kinetics_opt.default_diffuse_k0_upper_bound)
            self.assertEqual(actual_dx, mle_dx)
            self.assertEqual(len(actual_k0_lls_list), 1)
            self.assertEqual(actual_k0_lls_list[0][0], mle_k0)
            numpy.testing.assert_array_equal(actual_k0_lls_list[0][1], mle_lls)

    def test_full_probabilities_are_additive(self):
        t = 301.15
        p = wlc.default_p
        dx = 0.165e-9
        k0 = 0.039
        upper_bound_units = 400
        newtons_per_unit = 10e-12
        v0 = 400e-9
        l0 = 80e-9
        k = 0.06
        m0 = 3
        y_start_units0 = int(numpy.around(10e-12 / newtons_per_unit))
        protocol0 = kinetics_opt.Protocol(v0 / l0, k * l0, y_start_units0, m0)
        fp0 = kinetics_opt.full_probabilities(
            t, p, upper_bound_units, newtons_per_unit, [protocol0], dx, k0)
        self.assertEqual(fp0.shape, (1, upper_bound_units))
        numpy.testing.assert_allclose(fp0.sum(), 1)
        v1 = 8360e-9
        l1 = 10e-9
        k1 = 0.12
        m1 = 1
        y_start_units1 = int(numpy.around(200e-12 / newtons_per_unit))
        protocol1 = kinetics_opt.Protocol(v1 / l1, k * l1, y_start_units1, m1)
        fp1 = kinetics_opt.full_probabilities(
            t, p, upper_bound_units, newtons_per_unit, [protocol1], dx, k0)
        self.assertEqual(fp1.shape, (1, upper_bound_units))
        numpy.testing.assert_allclose(fp1.sum(), 1)
        expected = kinetics_opt.full_probabilities(
            t, p, upper_bound_units, newtons_per_unit, [protocol0, protocol1], dx, k0)
        self.assertEqual(expected.shape, (2, upper_bound_units))
        numpy.testing.assert_allclose(expected.sum(), 2)
        actual = numpy.concatenate([fp0, fp1])
        numpy.testing.assert_equal(actual, expected)


class SlowKineticsOptTest(unittest.TestCase):

    def test_estimate_dx_k0(self):
        t = 301.15
        p = wlc.default_p
        dx = 0.165e-9
        k0 = 0.039
        upper_bound_units = 400
        newtons_per_unit = 1e-12
        v = 400e-9
        l = 80e-9
        k = 0.06
        m = 3
        y_start_units = int(numpy.around(10e-12 / newtons_per_unit))
        protocol = kinetics_opt.Protocol(v / l, k * l, y_start_units, m)
        fp = kinetics_opt.full_probabilities(
            t, p, upper_bound_units, newtons_per_unit, [protocol], dx, k0)
        self.assertEqual(fp.shape, (1, upper_bound_units))
        n = 1000
        c = numpy.array(numpy.around(fp[0] * n), dtype=int)
        basic_peaks = sum(
            [[kinetics_opt.BasicPeak(protocol, i) for j in range(c[i])]
            for i in range(upper_bound_units)], [])
        mle_dx, mle_k0, mle_lls = kinetics_opt.estimate_dx_k0(
            t, p, upper_bound_units, newtons_per_unit, basic_peaks,
            0.01e-9, 1e-9, False)
        lls = kinetics_opt.log_likelihoods(
            t, p, dx, k0, upper_bound_units, newtons_per_unit, basic_peaks)
        logger.info(
            "dx(nm)=%r k0(s^{-1})=%r lls.sum()=%r lls(dx=%r(nm), k0=%r(s^{-1})=%r n=%d c.sum()=%d",
            mle_dx * 1e9, mle_k0, mle_lls.sum(), dx * 1e9, k0, lls.sum(), n, c.sum())
        numpy.testing.assert_allclose(mle_dx, dx, rtol=3e-3)

    def test_estimate_dx_k0_diffuse(self):
        t = 301.15
        p = wlc.default_p
        dx = 0.165e-9
        k0 = 0.039
        upper_bound_units = 400
        newtons_per_unit = 1e-12
        v = 400e-9
        l = 80e-9
        k = 0.06
        m = 3
        sigma_units = 20
        hsu = 60
        y_start_units = int(numpy.around(10e-12 / newtons_per_unit))
        diffuse_protocol = kinetics_opt.DiffuseProtocol(
            kinetics_opt.Protocol(v / l, k * l, y_start_units, m),
            sigma_units, hsu)
        for factory in (
            kinetics_opt.DiffusePeak.normal_from_basic_peak,
            kinetics_opt.DiffusePeak.truncated_normal_from_basic_peak,
            ):
            fp = kinetics_opt.full_probabilities_diffuse(
                t, p, upper_bound_units, newtons_per_unit, [diffuse_protocol], dx, k0,
                factory)
            self.assertEqual(fp.shape, (1, upper_bound_units))
            n = 1000
            c = numpy.array(numpy.around(fp[0] * n), dtype=int)
            diffuse_peaks = sum(
                [[factory(
                            kinetics_opt.BasicPeak(diffuse_protocol.protocol, i),
                            diffuse_protocol.sigma_units,
                            diffuse_protocol.hsu, upper_bound_units)
                  for j in range(c[i])]
                 for i in range(upper_bound_units)], [])
            mle_dx, mle_k0, mle_lls = kinetics_opt.estimate_dx_k0_diffuse(
                t, p, upper_bound_units, newtons_per_unit, diffuse_peaks,
                0.01e-9, 1e-9,
                kinetics_opt.default_diffuse_k0_lower_bound,
                kinetics_opt.default_diffuse_k0_upper_bound)
            lls = kinetics_opt.log_likelihoods_diffuse(
                t, p, dx, k0, upper_bound_units, newtons_per_unit, diffuse_peaks)
            logger.info(
                "dx(nm)=%r k0(s^{-1})=%r lls.sum()=%r lls(dx=%r(nm), k0=%r(s^{-1})=%r n=%d c.sum()=%d",
                mle_dx * 1e9, mle_k0, mle_lls.sum(), dx * 1e9, k0, lls.sum(), n, c.sum())
            numpy.testing.assert_allclose(mle_dx, dx, rtol=4e-2)


if __name__ == "__main__":
    suite = unittest.TestLoader().loadTestsFromTestCase(KineticsOptTest)
    unittest.TextTestRunner(verbosity=2).run(suite)
    suite = unittest.TestLoader().loadTestsFromTestCase(SlowKineticsOptTest)
    unittest.TextTestRunner(verbosity=2).run(suite)
